from tensorflow.python.keras import layers
import config
import tensorflow as tf

from model import Model
from utils import data_generator

if __name__ == '__main__':

    model = Model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=config.LR)
    for epoch in range(10):
        for index, (x, y, m) in enumerate(data_generator()):
            y_pred = model(x, True, m)

            loss = model.loss_fn(x, y, m)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(f"Epoch {epoch + 1}, Loss: {loss.numpy()}")

    print("Training finished.")
